import numpy as np
from utils.oss import download_file_from_oss, upload_file_to_oss
import trimesh
import pickle
from OCC.Core.TopoDS import TopoDS_Compound
from OCC.Core.BRep import BRep_Builder
import os
from OCC.Extend.DataExchange import read_step_file

def hole_sort(hole_list):
    if len(hole_list) == 1:
        return None
    elif len(hole_list) <= 4:
        target_axis = np.array(hole_list[1]) - np.array(hole_list[0])
    elif len(hole_list) == 6:
        # target_axis = np.array(start_target[i]['end'][0][1]) - np.array(start_target[i]['end'][0][0])
        target_axis_list = [
            (np.array(hole_list[0]) - np.array(hole_list[k])) for k in
            range(1, 6)]
        # print(start_axis_list)
        temp = target_axis_list[0]
        for j in range(5):
            if np.linalg.norm(target_axis_list[j]) < np.linalg.norm(temp):
                temp = target_axis_list[j]
        target_axis = temp
    target_axis = target_axis / np.linalg.norm(target_axis)
    return target_axis


class PointsJsonParser():
    def __init__(self,json_file):
        self.start_target = json_file['begin_end']
        self.work_path=json_file['WorkPath']
        self.rps_list=json_file['Rps']
        self.pickle = json_file['pickle']
        self.local_work_path = os.path.dirname(__file__)+'temp_new.stl'
        self.local_pickle_path = os.path.dirname(__file__)+'clamp_new.pickle'
        if "adjust_times" in json_file:
            self.adjust_times=json_file['adjust_times']
        else:
            self.adjust_times = 1
        # self.adjust_times = json_file['adjust_times']

    def set_local_work_path(self,work_path):
        self.local_work_path = work_path

    def set_local_pickle_path(self,pickle_path):
        self.local_pickle_path = pickle_path

    def pickle2compound(self):
        if isinstance(self.pickle,str):
            if self.pickle.startswith('DO-Design_Dev'):
                download_file_from_oss(oss_file_name=self.pickle,
                                       local_file_name=self.local_pickle_path)
            else:
                self.local_pickle_path = self.pickle
        p2 = open(self.local_pickle_path, 'rb')
        group_list = pickle.load(p2)
        p2.close()
        compound = TopoDS_Compound()
        B = BRep_Builder()
        B.MakeCompound(compound)
        for group in group_list:
            group.compound()
            B.Add(compound, group.compound())
        return compound

    def load_stp(self):
        if isinstance(self.pickle,str):
            if self.pickle.startswith('DO-Design_Dev'):
                download_file_from_oss(oss_file_name=self.pickle,
                                       local_file_name=self.local_pickle_path)
            else:
                self.local_pickle_path = self.pickle
        compound = read_step_file(self.local_pickle_path)
        return compound

    def get_work_mesh(self):
        if isinstance(self.work_path, str):
            if self.work_path.startswith('DO-Design_Dev'):
                download_file_from_oss(oss_file_name=self.work_path,
                                       local_file_name=self.local_work_path)
            else:
                self.local_work_path = self.work_path
        work_mesh = trimesh.load(self.local_work_path)
        return work_mesh

    def get_points_pair(self):
        points_pair = []
        start_target = self.start_target
        if len(self.rps_list)==0:
            for i in range(len(start_target)):
                start = np.mean(start_target[i]['begin'][0], axis=0)
                start_axis = hole_sort(start_target[i]['begin'][0])
                if isinstance(start_target[i]['begin'][1][0],list):
                    start_normal = np.array(start_target[i]['begin'][1][0])
                else:
                    start_normal = np.array(start_target[i]['begin'][1])
                start_normal = start_normal / np.linalg.norm(start_normal)

                end_list = start_target[i]['end']
                if len(end_list) == 1:
                    target = np.mean(start_target[i]['end'][0][0], axis=0)
                    target_axis = hole_sort(start_target[i]['end'][0][0])
                    if isinstance(start_target[i]['end'][0][1][0], list):
                        target_normal = np.array(start_target[i]['end'][0][1][0])
                    else:
                        target_normal = np.array(start_target[i]['end'][0][1])
                    target_normal = target_normal/np.linalg.norm(target_normal)
                elif len(end_list) > 1:
                    target = []
                    target_axis = []
                    target_normal = []
                    for j in range(len(end_list)):
                        target.append(np.mean(end_list[j][0], axis=0))
                        target_axis_temp = hole_sort(end_list[j][0])
                        target_axis.append(target_axis_temp)
                        if isinstance(end_list[j][1][0],list):
                            target_normal.append(end_list[j][1][0])
                        else:
                            target_normal.append(end_list[j][1])
                # target = np.mean(start_target[i]['end'][0], axis=0)
                # target_axis = hole_sort(start_target[i]['end'][0])
                    # target_normal = np.array(start_target[i]['end'][1][0])

                pair = {
                    "start": start,
                    "target": target,
                    "start_axis": start_axis,
                    "target_axis": target_axis,
                    "start_normal": start_normal,
                    "target_normal": target_normal,
                    "adjust_times":self.adjust_times,
                    "rps_type": None
                }
                points_pair.append(pair)
        elif len(self.rps_list) == 1:
            rps_type = self.rps_list[0]['Type']
            point = self.rps_list[0]['Point']
            for i in range(len(start_target)):
                start = np.mean(start_target[i]['begin'][0], axis=0)
                start_axis = hole_sort(start_target[i]['begin'][0])
                start_normal = np.array(start_target[i]['begin'][1])
                start_normal = start_normal / np.linalg.norm(start_normal)
                if "end" not in start_target[i] or start_target[i]['end'] is None:
                    target = np.array(point)
                    target_axis = None
                    target_normal = -start_normal
                else:
                    end_list = start_target[i]['end']
                    if len(end_list) == 1:
                        target = np.mean(start_target[i]['end'][0][0], axis=0)
                        target_axis = hole_sort(start_target[i]['end'][0][0])
                        target_normal = np.array(start_target[i]['end'][0][1])
                        target_normal = target_normal/np.linalg.norm(target_normal)
                    elif len(end_list) > 1:
                        target = []
                        target_axis = []
                        target_normal = []
                        for j in range(len(end_list)):
                            target.append(np.mean(end_list[j][0], axis=0))
                            target_axis_temp = hole_sort(end_list[j][0])
                            target_axis.append(target_axis_temp)
                            target_normal.append(end_list[j][1])
                    # target = np.mean(start_target[i]['end'][0], axis=0)
                    # target_axis = hole_sort(start_target[i]['end'][0])
                    # target_normal = np.array(start_target[i]['end'][1][0])

                pair = {
                    "start": start,
                    "target": target,
                    "start_axis": start_axis,
                    "target_axis": target_axis,
                    "start_normal": start_normal,
                    "target_normal": target_normal,
                    "adjust_times":self.adjust_times,
                    "rps_type": None
                }
                if rps_type == "伸缩销":
                    pair["rps_type"] = "伸缩销"
                points_pair.append(pair)
        elif len(self.rps_list) > 1:
            for i in range(len(start_target)):
                start = np.mean(start_target[i]['begin'][0], axis=0)
                start_axis = hole_sort(start_target[i]['begin'][0])
                start_normal = np.array(start_target[i]['begin'][1])
                start_normal=start_normal/np.linalg.norm(start_normal)
                if "end" not in start_target[i] or start_target[i]['end'] is None:
                    point = [self.rps_list[i]['Point'] for i in range(len(self.rps_list))]
                    target = np.array(point)
                    target_axis = None
                    target_normal = -start_normal

                elif len(start_target[i]['end'][0])==1:
                    target=[np.array(start_target[i]['end'][k][0]) for k in range(len(start_target[i]['end']))]
                    target_axis = None
                    target_normal = [-start_normal for k in range(len(start_target[i]['end']))]
                    rps_type = None
                else:
                    end_list = start_target[i]['end']
                    if len(end_list) == 1:
                        target = np.mean(start_target[i]['end'][0], axis=0)
                        target_axis = hole_sort(start_target[i]['end'][0])
                        target_normal = np.array(start_target[i]['end'][1])/np.linalg.norm(np.array(start_target[i]['end'][1]))
                    elif len(end_list) > 1:
                        target = []
                        target_axis = []
                        target_normal = []
                        for j in range(len(end_list)):
                            target.append(np.mean(end_list[j][0], axis=0))
                            target_axis_temp = hole_sort(end_list[j][0])
                            target_axis.append(target_axis_temp)
                            target_normal.append(np.array(end_list[j][1])/np.linalg.norm(np.array(end_list[j][1])))
                    # target = np.mean(start_target[i]['end'][0], axis=0)
                    # target_axis = hole_sort(start_target[i]['end'][0])
                    # target_normal = np.array(start_target[i]['end'][1])

                pair = {
                    "start": start,
                    "target": target,
                    "start_axis": start_axis,
                    "target_axis": target_axis,
                    "start_normal": start_normal,
                    "target_normal": target_normal,
                    "rps_type": None
                }
                # if rps_type == "伸缩销":
                #     pair["rps_type"] = "伸缩销"
                points_pair.append(pair)
        return points_pair

if __name__ == '__main__':
    import utils.json_tools as json_tool
    json = json_tool.read_json('sb2double_RPS.json')
    parser = PointsJsonParser(json)

    print(parser.pickle2compound())
    from OCC.Display.SimpleGui import init_display

    display, start_display, add_menu, add_function_to_menu = init_display("wx")



    display.DisplayShape(parser.pickle2compound(), update=True)
    start_display()
